#include "demo_utils.hpp"

#include <atomic>
#include <iostream>
#include <optional>
#include <thread>
#include <vector>

namespace atomic_examples {
namespace {

struct Node {
  int value;   ///< 节点保存的数值。
  Node* next;  ///< 指向下一节点的指针。
};

/**
 * @brief 使用 CAS（比较交换）构建的简单无锁栈。
 *
 * 设计说明：
 *  - `head_` 为栈顶指针，使用 `std::atomic<Node*>` 存储。
 *  - push/pop 操作都基于 `compare_exchange_weak` 循环完成，确保线程安全。
 *  - 适用于单机内低争用场景的快速栈实现。
 *
 * ⚠️ 注意：
 *  - 未处理 ABA 问题，如需更严格保证，应引入版本号或指针标记。
 *  - 堆上分配的节点需要在 pop 后及时释放，避免内存泄漏。
 */
class LockFreeStack {
public:
  /**
   * @brief 压入一个整数。
   *
   * @param value 待压入的整数值。
   */
  void push(int value) {
    Node* new_node = new Node{value, nullptr};
    Node* old_head = head_.load(std::memory_order_relaxed);

    while (true) {
      new_node->next = old_head;  ///< 先链接到旧栈顶节点。
      if (head_.compare_exchange_weak(
              old_head,
              new_node,
              std::memory_order_release,
              std::memory_order_relaxed)) {
        return;  ///< ✅ 成功将 new_node 发布给其他线程。
      }
      // compare_exchange_weak 失败时 old_head 会被更新为最新值，循环继续。
    }
  }

  /**
   * @brief 弹出栈顶元素。
   *
   * @return 若弹出成功，返回数据；若栈为空，返回 std::nullopt。
   */
  std::optional<int> pop() {
    Node* old_head = head_.load(std::memory_order_acquire);

    while (old_head) {
      Node* next = old_head->next;
      if (head_.compare_exchange_weak(
              old_head,
              next,
              std::memory_order_acquire,
              std::memory_order_relaxed)) {
        int value = old_head->value;
        delete old_head;  ///< ✅ 释放节点内存。
        return value;
      }
      // 若 CAS 失败，old_head 已更新为新的栈顶，继续循环尝试。
    }
    return std::nullopt;
  }

private:
  std::atomic<Node*> head_{nullptr}; ///< 栈顶指针，使用原子指针以避免数据竞争。
};

/**
 * @brief 演示多线程生产者向无锁栈压入元素，并最终全部弹出。
 */
void runLockFreeStackDemo() {
  printHeader("Lock-Free Stack Example");

  LockFreeStack stack;
  constexpr int kProducers = 4;
  constexpr int kPerProducer = 2'500;

  std::vector<std::thread> producers;
  producers.reserve(kProducers);

  for (int p = 0; p < kProducers; ++p) {
    producers.emplace_back([&, p]() {
      for (int i = 0; i < kPerProducer; ++i) {
        stack.push(p * kPerProducer + i);
      }
    });
  }
  for (auto& t : producers) {
    t.join();
  }

  int pop_count = 0;
  while (stack.pop().has_value()) {
    ++pop_count;
  }

  std::cout << "Total pushed: " << (kProducers * kPerProducer)
            << ", Total popped: " << pop_count << '\n';
}

} // namespace
} // namespace atomic_examples

int main() {
  atomic_examples::runLockFreeStackDemo();
  return 0;
}